package cc.xpcas.raft.role;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import cc.xpcas.raft.LogEntry;
import cc.xpcas.raft.context.Context;
import cc.xpcas.raft.context.Peer;
import cc.xpcas.raft.message.AppendEntriesRequest;
import cc.xpcas.raft.message.AppendEntriesResponse;
import cc.xpcas.raft.message.MessageContext;
import lombok.Data;

/**
 * @author xp
 */
public class Leader extends Role {

    private static final Logger LOG = LoggerFactory.getLogger(Leader.class);

    @Data
    private static class PeerState {

        private final Peer peer;

        private long lastResponseTimestamp = 0L;

        private long lastLogId = 0L;

        private long lastCommittedLogId = 0L;

        public PeerState(Peer peer) {
            this.peer = peer;
        }
    }

    private final Map<Integer, PeerState> stateMap = new ConcurrentHashMap<>();

    private final AtomicBoolean timeoutFlag = new AtomicBoolean(false);

    private Future future;

    public Leader(Context context) {
        super(context);
    }

    /**
     * in mainline
     */
    @Override
    public void start() {
        LOG.info(String.format("%s: leader start", getConfiguration().getServerId()));

        getContext().getPeers().forEach((serverId, peer) -> {
            PeerState state = new PeerState(peer);
            stateMap.put(serverId, state);
        });
        tryAppendingEntriesToAll(0); // heartbeat
        resumeHeartbeatCountdown();
    }

    /**
     * in mainline
     */
    @Override
    public void end() throws InterruptedException {
        stopHeartbeatCountdown();
    }

    /**
     * in mainline
     */
    private void stopHeartbeatCountdown() {
        timeoutFlag.set(true);
        if (future != null) {
            future.cancel(false);
        }
    }

    /**
     * in mainline
     * in scheduled
     */
    private void resumeHeartbeatCountdown() {
        timeoutFlag.set(false);
        future = getContext().schedule(this::heartbeatHandler, getConfiguration().getHeartbeatCheckInterval());
    }

    /**
     * in mainline
     */
    private void appendEntries(PeerState state, int bulkLimit) {
        AppendEntriesRequest request = new AppendEntriesRequest();
        request.setServerId(getConfiguration().getServerId());
        request.setTerm(getContext().getTerm());
        request.setLastCommittedLogId(getContext().getLastCommittedLogId());

        LogEntry lastPreparedLog = getContext().getLogEntry(state.getLastLogId());
        if (lastPreparedLog != null) {
            request.setPreviousLogId(lastPreparedLog.getId());
            request.setPreviousLogTerm(lastPreparedLog.getTerm());
        } else {
            request.setPreviousLogId(0);
            request.setPreviousLogTerm(0);
        }

        List<LogEntry> entries = new ArrayList<>(bulkLimit);
        if (bulkLimit > 0) {
            Iterator<LogEntry> iterator = getContext().getLogEntriesSince(request.getPreviousLogId() + 1);
            for (int i = bulkLimit; i > 0 && iterator.hasNext(); --i) {
                entries.add(iterator.next());
            }
        }
        request.setEntries(entries);

        state.getPeer().send(request);
        state.setLastResponseTimestamp(System.currentTimeMillis());
    }

    /**
     * in mainline
     */
    private void tryAppendingEntriesToAll(int bulkLimit) {
        if (!getContext().isLeader()) return;

        stateMap.values().forEach(state ->
                tryAppendingEntriesTo(state, bulkLimit));
    }

    /**
     * in mainline
     */
    private void tryAppendingEntriesTo(PeerState state, int bulkLimit) {
        if (!isHeartbeatRequired(state.getLastResponseTimestamp())) return;

        appendEntries(state, bulkLimit);
    }

    private boolean isHeartbeatRequired(long lastResponseTimestamp) {
        long diff = System.currentTimeMillis() - lastResponseTimestamp;
        return diff >= getConfiguration().getHeartbeatInterval();
    }

    private boolean isNewLeader(long term) {
        /*
         * -1 new leader
         *  0 error for request
         *  1 old leader
         */
        return getContext().catchUpTerm(term) < 0;
    }

    /**
     * in mainline
     */
    @Override
    public boolean handleAppendEntriesResponse(MessageContext messageContext) throws InterruptedException {
        AppendEntriesResponse message = messageContext.getMessage();
        if (isNewLeader(message.getCurrentTerm())) {
            // we're follower now
            return true;
        }
        int serverId = message.getServerId();
        PeerState state = stateMap.get(serverId);
        assert state != null;

        state.setLastLogId(message.getLastLogId());
        state.setLastCommittedLogId(message.getLastCommittedLogId());
        state.setLastResponseTimestamp(System.currentTimeMillis());

        return true;
    }

    /**
     * in mainline
     */
    @Override
    public boolean handleAppendEntriesRequest(MessageContext messageContext) throws InterruptedException {
        AppendEntriesRequest message = messageContext.getMessage();
        if (isNewLeader(message.getTerm())) {
            // we're follower now
            return false;
        }
        return true;
    }

    /**
     * in scheduled
     */
    private void heartbeatHandler() {
        if (!timeoutFlag.compareAndSet(false, true)) {
            // already closed
            return;
        }
        //getContext().runOnMainline(() -> tryAppendingEntriesToAll(getConfiguration().getAppendEntriesLimit()));
        resumeHeartbeatCountdown();
    }
}
